import sqlite3
import re
from typing import Optional

class CVEDatabase:
    def __init__(self, db_name='cve.db'):
        self.conn = sqlite3.connect(db_name)
        self._create_tables()
        
    def _create_tables(self):
        """创建基础表结构"""
        self.conn.execute('''CREATE TABLE IF NOT EXISTS CVE (
                        cve_id INTEGER PRIMARY KEY,
                        name TEXT UNIQUE,
                        intro_mainline_id TEXT,
                        fixed_mainline_id TEXT
                      )''')

        self.conn.execute('''CREATE TABLE IF NOT EXISTS OS_Version (
                        os_id INTEGER PRIMARY KEY,
                        version_name TEXT UNIQUE,
                        version_type TEXT CHECK(version_type IN ('stable', 'openEuler')),
                        parent TEXT,
                        CHECK (
                            (version_type = 'openEuler' AND parent IS NOT NULL) OR
                            (version_type = 'stable' AND parent IS NOT NULL)
                        )
                        )''')

        # 创建触发器确保 openEuler 版本关联的父版本存在且为 stable 类型
        self.conn.execute('''CREATE TRIGGER IF NOT EXISTS validate_openeuler_parent_insert
                        BEFORE INSERT ON OS_Version
                        FOR EACH ROW
                        WHEN NEW.version_type = 'openEuler'
                        BEGIN
                            SELECT RAISE(ABORT, 'Parent must be an existing stable version')
                            WHERE NOT EXISTS (
                                SELECT 1 FROM OS_Version 
                                WHERE version_name = NEW.parent 
                                AND version_type = 'stable'
                            );
                        END;''')

        self.conn.execute('''CREATE TRIGGER IF NOT EXISTS validate_openeuler_parent_update
                        BEFORE UPDATE ON OS_Version
                        FOR EACH ROW
                        WHEN NEW.version_type = 'openEuler'
                        BEGIN
                            SELECT RAISE(ABORT, 'Parent must be an existing stable version')
                            WHERE NOT EXISTS (
                                SELECT 1 FROM OS_Version 
                                WHERE version_name = NEW.parent 
                                AND version_type = 'stable'
                            );
                        END;''')

        self.conn.execute('''CREATE TABLE IF NOT EXISTS CVE_OS_Link (
                        cve_id INTEGER,
                        os_id INTEGER,
                        commit_id TEXT,
                        PRIMARY KEY (cve_id, os_id),
                        FOREIGN KEY (cve_id) REFERENCES CVE(cve_id),
                        FOREIGN KEY (os_id) REFERENCES OS_Version(os_id)
                        )''')
        self.conn.commit()
        
    def check_os_version_is_openeuler(self, version_name):
        """
        检查指定的操作系统版本是否为openEuler类型
        :param version_name: 要检查的操作系统版本名称
        :return: 
            - True: 是openEuler版本
            - False: 不是openEuler版本或版本不存在
        """
        try:
            cursor = self.conn.execute(
                'SELECT version_type FROM OS_Version WHERE version_name = ?',
                (version_name,)
            )
            result = cursor.fetchone()
            # 如果记录存在且类型匹配
            return result is not None and result[0] == "openEuler"
        except sqlite3.DatabaseError as e:
            print(f"数据库查询错误: {e}")
            return False

    def get_os_versions(self, type: str = "all") -> list:
        """
        获取操作系统版本（严格类型限制）
        :param type: 过滤类型
            - "all": 返回所有记录（默认）
            - "openEuler": 仅返回openEuler类型
            - "stable": 仅返回stable类型
        :return: 包含完整信息的字典列表
        :raises ValueError: 当传入无效type参数时
        """
        # 严格参数验证
        if type not in ("all", "openEuler", "stable"):
            raise ValueError(f"无效类型参数: {type}，只允许 'all'/'openEuler'/'stable'")

        try:
            query = 'SELECT version_name FROM OS_Version'
            params = ()
            if type != "all":
                query += " WHERE version_type = ?"
                params = (type,)
            query += " ORDER BY version_name"
            # 执行查询
            cursor = self.conn.execute(query, params)
            return [row[0] for row in cursor.fetchall()]

        except sqlite3.DatabaseError as e:
            print(f"数据库查询失败: {e}")
            return []
        
    def get_os_parents(self, os_version: str) -> Optional[str]:
        """直接返回 parent 字段的值"""
        try:
            cursor = self.conn.cursor()
            cursor.execute('SELECT parent FROM OS_Version WHERE version_name = ?', (os_version,))
            result = cursor.fetchone()
            return result[0] if result else None
        except sqlite3.Error as e:
            print(f"Database error: {str(e)}")
            return None
    
    def get_all_cves(self) -> list[dict]:
        """
        获取数据库中所有CVE记录
        :return: 包含CVE详细信息的字典列, 结构为:
            [{
                "cve_id": int,
                "name": str,
                "intro_mainline_id": str,
                "fixed_mainline_id": str
            }, ...]
        """
        try:
            cursor = self.conn.execute(
                '''SELECT cve_id, name, intro_mainline_id, fixed_mainline_id FROM CVE ORDER BY cve_id'''
            )
            return [{
                    "cve_id": row[0],
                    "name": row[1],
                    "intro_mainline_id": row[2],
                    "fixed_mainline_id": row[3]
                }
                for row in cursor.fetchall()
            ]
        except sqlite3.DatabaseError as e:
            print(f"获取CVE列表失败: {e}")
            return []

    def check_fixed_commit_id(self, cve_name: str, os_version: str) -> Optional[str]:
        """
        获取指定CVE在特定操作系统版本下的修复commit标识
        :param cve_name: CVE名称
        :param os_version: 操作系统版本名称
        :return: 
            - str: 找到的commit链接标识
            - None: 未找到记录或发生错误
        """
        try:
            # 原子性获取所有必要ID
            cursor = self.conn.execute('''
                SELECT l.commit_id
                FROM CVE_OS_Link l
                JOIN CVE c ON l.cve_id = c.cve_id
                JOIN OS_Version o ON l.os_id = o.os_id
                WHERE c.name = ? AND o.version_name = ?
            ''', (cve_name, os_version))
            
            result = cursor.fetchone()
            return result[0] if result else None

        except sqlite3.DatabaseError as e:
            print(f"数据库查询失败: {e}")
            return None

    def check_intro_commit_id_mainline(self, cve_name: str) -> Optional[str]:
        """
        获取指定CVE的引入commit标识
        :param cve_name: CVE名称
        :return: 
            - str: 引入commit标识
            - None: CVE不存在或查询失败
        """
        try:
            cursor = self.conn.execute("SELECT intro_mainline_id FROM CVE WHERE name = ?", (cve_name,))
            result = cursor.fetchone()
            return result[0] if result else None
        except sqlite3.DatabaseError as e:
            print(f"查询引入commit失败: {e}")
            return None
        
    def check_fixed_commit_id_mainline(self, cve_name: str) -> Optional[str]:
        """
        获取指定CVE在mainline分支的修复commit标识 (mainline)
        :param cve_name: CVE名称
        :return: 
            - str: 有效修复commit标识
            - None: 无修复记录/CVE不存在/数据异常
        """
        try:
            cursor = self.conn.execute("SELECT fixed_mainline_id FROM CVE WHERE name = ?", (cve_name,))
            result = cursor.fetchone()
            return result[0] if result else None
        except sqlite3.DatabaseError as e:
            print(f"数据库查询失败: {e}")
            return None

    def check_os_version_is_stable_version(self, version_name: str) -> Optional[str]:
        """
        检查并返回匹配的stable基准版本名称
        :param version_name: 要检查的版本名称
        :return: 
            - str: 匹配的stable基准版本名称
            - None: 无匹配版本或发生错误
        """
        def parse_version(ver: str) -> list[int]:
            """解析版本字符串为数字序列（忽略非数字后缀）"""
            parts = []
            for part in ver.split('.'):
                match = re.match(r'^(\d+)', part)
                if match:
                    parts.append(int(match.group(1)))
                else:
                    break
            return parts

        try:
            # 获取所有stable基准版本并解析
            cursor = self.conn.execute(
                "SELECT version_name FROM OS_Version "
                "WHERE version_type = 'stable'"
            )
            
            # 存储格式：{"name": "OLK-5.10", "parts": [5,10]}
            base_versions = []
            for (db_version,) in cursor.fetchall():
                version_part = self.get_os_parents(db_version)
                parts = parse_version(version_part)
                if parts:
                    base_versions.append({
                        "name": db_version,
                        "parts": parts
                    })

            # 按版本段长度降序排序（优先匹配更长更精确的版本）
            base_versions.sort(key=lambda x: len(x["parts"]), reverse=True)

            # 解析输入版本
            input_parts = parse_version(version_name)
            if not input_parts:
                return None

            # 寻找最长匹配
            for base in base_versions:
                base_parts = base["parts"]
                # 输入版本需至少包含基准版本的所有段
                if len(input_parts) < len(base_parts):
                    continue
                # 严格前缀匹配
                if input_parts[:len(base_parts)] == base_parts:
                    return base["name"]

            return None

        except sqlite3.DatabaseError as e:
            print(f"数据库查询失败: {e}")
            return None
        except Exception as e:
            print(f"版本解析异常: {e}")
            return None

    def add_os_version(self, version_name: str, version_type: str, parent: str):
        """
        添加新的操作系统版本
        :param version_name: 版本名称 (如 "openEuler-22.03-LTS")
        :param version_type: 系统类型 (stable/openEuler)
        :param parent: 父版本名称（openEuler类型必须为存在的stable版本名称）
        """
        cursor = self.conn.cursor()
        try:
            # 参数基础验证
            if version_type not in ("stable", "openEuler"):
                raise ValueError("version_type 必须为 stable 或 openEuler")
            if not parent:
                raise ValueError("parent 参数不能为空")

            # 直接插入 parent 字符串（触发器会验证 openEuler 类型的父版本有效性）
            cursor.execute('''
                INSERT INTO OS_Version 
                (version_name, version_type, parent)
                VALUES (?, ?, ?)
            ''', (version_name, version_type, parent))
                
            self.conn.commit()
            print(f"成功添加 {version_type} 版本: {version_name}")

        except sqlite3.IntegrityError as e:
            self.conn.rollback()
            if "UNIQUE constraint failed" in str(e):
                print(f"版本名称 '{version_name}' 已存在")
            elif "CHECK constraint failed" in str(e):
                print(f"父版本参数不符合约束条件")
            else:
                print(f"数据完整性错误: {e}")
        except Exception as e:
            self.conn.rollback()
            raise RuntimeError(f"添加操作系统版本失败: {str(e)}") from e
        
    def add_cve(self, cve_name: str, intro_mainline_id: str, fixed_mainline_id: str) -> Optional[int]:
        """
        添加或更新 CVE 到数据库
        返回新插入或更新的 cve_id, 若失败返回 None
        """
        # 参数校验
        if not re.match(r'^CVE-\d{4}-\d+$', cve_name):
            raise ValueError(f"Invalid CVE name: {cve_name}. Expected format: CVE-YYYY-NNNN")
        try:
            cur = self.conn.cursor()
            # 使用 UPSERT 语法插入或更新 CVE 记录
            cur.execute('''
                INSERT INTO CVE (name, intro_mainline_id, fixed_mainline_id)
                VALUES (?, ?, ?)
                ON CONFLICT(name) DO UPDATE SET
                    intro_mainline_id = excluded.intro_mainline_id,
                    fixed_mainline_id = excluded.fixed_mainline_id
            ''', (cve_name, intro_mainline_id, fixed_mainline_id))
            # 获取受影响的行数，如果是更新操作，则返回现有的 cve_id
            self.conn.commit()
            if cur.rowcount == 0:
                # 新插入的记录
                cve_id = cur.lastrowid
            else:
                # 更新的记录，获取现有的 cve_id
                cur.execute('SELECT cve_id FROM CVE WHERE name = ?', (cve_name,))
                cve_id = cur.fetchone()[0]
            return cve_id
        except sqlite3.IntegrityError as e:
            self.conn.rollback()
            # 解析具体违反的约束
            error_msg = str(e)
            if "CVE.name" in error_msg:
                print(f"CVE {cve_name} 已存在")
            elif "CVE.intro_mainline_id" in error_msg:
                print(f"主线版本 {intro_mainline_id} 已被其他 CVE 使用")
            elif "CVE.fixed_mainline_id" in error_msg:
                print(f"主线版本 {fixed_mainline_id} 已被其他 CVE 使用")
            else:
                raise RuntimeError(f"数据库错误: {error_msg}") from None
        except Exception as e:
            self.conn.rollback()
            raise RuntimeError(f"操作失败: {str(e)}") from None

    def add_cve_commit(self, cve_name: str, os_version: str, commit_id: str) -> bool:
        """
        更新/插入CVE与操作系统版本的关联修复commit
        :param cve_name: CVE名称
        :param os_version: 操作系统版本名称
        :param commit_id: 修复commit链接标识
        :return: 
            - True: 操作成功
            - False: 操作失败 (CVE/版本不存在或数据库错误)
        """
        try:
            # 获取CVE ID
            cve_row = self.conn.execute("SELECT cve_id FROM CVE WHERE name = ?", (cve_name,)).fetchone()
            if not cve_row:
                print(f"错误: CVE '{cve_name}' 不存在")
                return False
            cve_id = cve_row[0]

            # 获取OS版本ID
            os_row = self.conn.execute("SELECT os_id FROM OS_Version WHERE version_name = ?", (os_version,)).fetchone()
            if not os_row:
                print(f"错误：操作系统版本 '{os_version}' 不存在")
                return False
            os_id = os_row[0]

            # 执行UPSERT操作
            self.conn.execute(
                """INSERT INTO CVE_OS_Link (cve_id, os_id, commit_id)
                VALUES (?, ?, ?)
                ON CONFLICT(cve_id, os_id) 
                DO UPDATE SET commit_id = excluded.commit_id""",
                (cve_id, os_id, commit_id)
            )
            self.conn.commit()
            return True

        except sqlite3.IntegrityError as e:
            print(f"数据完整性错误: {e}")
            self.conn.rollback()
            return False
        except sqlite3.DatabaseError as e:
            print(f"数据库操作失败: {e}")
            self.conn.rollback()
            return False
        
    def update_os_version_name(self, old_version_name: str, new_version_name: str) -> bool:
        """
        更新操作系统版本名称（级联更新 openEuler 子版本的 parent 字段）
        """
        try:
            with self.conn:
                cursor = self.conn.cursor()
                
                # 检查新名称是否已存在
                cursor.execute("SELECT 1 FROM OS_Version WHERE version_name = ?", (new_version_name,))
                if cursor.fetchone():
                    print(f"错误: 版本名称 '{new_version_name}' 已存在")
                    return False

                # 获取旧版本信息
                cursor.execute("SELECT version_type, parent FROM OS_Version WHERE version_name = ?", (old_version_name,))
                record = cursor.fetchone()
                if not record:
                    print(f"错误: 找不到版本 '{old_version_name}'")
                    return False
                old_version_type, old_parent = record

                # 更新自身版本名称
                cursor.execute(
                    "UPDATE OS_Version SET version_name = ? WHERE version_name = ?",
                    (new_version_name, old_version_name)
                )

                # 如果是 stable 版本，级联更新所有 openEuler 子版本的 parent 字段
                if old_version_type == "stable":
                    cursor.execute(
                        "UPDATE OS_Version SET parent = ? WHERE parent = ? AND version_type = 'openEuler'",
                        (new_version_name, old_version_name)
                    )

                return True

        except sqlite3.Error as e:
            print(f"数据库错误: {str(e)}")
            return False
        
    def close(self):
        self.conn.close()
